import torch
import torch.nn as nn
import torch.optim as optim
import pandas as pd
import numpy as np
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import StandardScaler
from model import LSTMModel

# 训练函数
def train_model(model, train_loader, criterion, optimizer, num_epochs=50):
    model.train()
    for epoch in range(num_epochs):
        for data, labels in train_loader:
            data, labels = data.to(device), labels.to(device)

            # 前向传播
            outputs = model(data)
            loss = criterion(outputs, labels)

            # 反向传播和优化
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()

        print(f'Epoch [{epoch+1}/{num_epochs}], Loss: {loss.item():.4f}')
    
    print("训练完成，开始保存模型...")
    torch.save(model.state_dict(), "lstm_model.pth")
    print("模型已保存为 lstm_model.pth")


# 数据处理和加载
def load_data(csv_file):
    df = pd.read_csv(csv_file)

    # 提取输入特征 (A1, A2, ..., B120) 和标签 (match_label)
    X = df.drop(columns=["match_label"]).values
    y = df["match_label"].apply(lambda x: 1 if x == "team1" else 0).values

    # 标准化特征
    scaler = StandardScaler()
    X = scaler.fit_transform(X)

    # 拆分训练集和测试集
    X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)

    # 将数据转为 PyTorch tensor
    X_train = torch.tensor(X_train, dtype=torch.float32).reshape(-1, 1, 360)
    y_train = torch.tensor(y_train, dtype=torch.long)

    train_dataset = torch.utils.data.TensorDataset(X_train, y_train)
    train_loader = torch.utils.data.DataLoader(dataset=train_dataset, batch_size=32, shuffle=True)

    return train_loader


if __name__ == "__main__":
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

    # 加载数据
    csv_file = 'matches_odds.csv'
    train_loader = load_data(csv_file)

    # 初始化模型、损失函数和优化器
    model = LSTMModel(input_size=360, hidden_size=128, num_layers=2, output_size=2).to(device)
    criterion = nn.CrossEntropyLoss()
    optimizer = optim.Adam(model.parameters(), lr=0.001)

    # 训练模型
    train_model(model, train_loader, criterion, optimizer, num_epochs=50)
